#!/usr/bin/python3
# -*- coding: utf-8 -*-
import re
import sys

from setuptools import setup, find_packages

global dev_depends, prod_depends

dev_depends = [
    'tensorflow-2.3.0'
]

prod_depends = [
    'ujson>=1.35',
    'numpy>=1.16.4',
    'PyYAML>=5.1.2',
    'transformers-3.0.2'
]


def get_deps_by_level(level):
    global dev_depends, prod_depends
    if level == 'prod':
        final_depends = prod_depends
    elif level == 'all':
        final_depends = set(dev_depends).union(prod_depends)
    else:
        final_depends = set()
    final_depends = list(final_depends)
    return final_depends


set_up_params = {
    'name': 'DistillBert',
    'version': '1.0',
    'author': 'beisen',
    'description': 'distill albert into bilistm model',
    'packages': find_packages(),
    'install_requires': [],
    'package_data': {
        'data': ['*.txt', '*.json', '*.yml', '*.pkl', '*.model'],
        'env': ['*.txt'],
        '': ['*.yml']
    }
}

reg_level = re.compile(r'\-\-level[ =](\w+)')

if __name__ == '__main__':
    call_args = ' '.join(sys.argv)
    mat_level = reg_level.search(call_args)
    if mat_level:
        level = mat_level.group(1)
        start_id = next((i for i, c in enumerate(sys.argv) if c == '--level'), -1)
        if start_id > 0:
            sys.argv = sys.argv[:start_id] + sys.argv[start_id + 2:]
    else:
        level = None
    require_deps = get_deps_by_level(level)
    set_up_params['install_requires'] = require_deps
    setup(**set_up_params)
